
import sys

class PadSequence:
    LEFT = 1
    RIGHT = 2
    NONE = 0
    BOTH = 3

def readSequence(filename, spos, epos, seqName=None, returnAsFasta=False, fastaLineLength=60, pad=PadSequence.BOTH):
    '''
    Read a range of positions from a FASTA file, numbering from 1.
    Return the specified bases either as a string or a new FASTA.
    Depending on the padding policy, add enough N's to the beginning
    and/or the end of the sequence to fill in the specified range.
    '''

    inSeq = False
    currSeq = None
    padLeft = pad & PadSequence.LEFT
    padRight = pad & PadSequence.RIGHT
    firstPos = spos
    lastPos = epos
    ret = ""

    if padLeft:
        while spos<=0:
            spos += 1
            ret += "N"
    else:
        if spos<=0:
            spos = 1
            firstPos = 1
    
    for line in open(filename):
        if len(line)>0 and line[0]=='>':
            pos = 1
            currSeq = line[1:].strip()
            inSeq = (seqName==None or currSeq==seqName)
        elif inSeq:
            seq = line.strip().upper()
            if pos > epos:
                break
            elif pos + len(seq) <= spos:
                pos += len(seq)
            else:
                # pos + len(seq) >= spos
                # pos <= epos
                ret += seq[max(0, spos-pos) : epos-pos+1]
                pos += len(seq)
    
    if padRight:
        while pos<=epos:
            pos += 1
            ret += "N"
    else:
        if pos<=epos:
            lastPos = pos-1
            
    if not returnAsFasta:
        return ret

    rng = "[%d,%d]" % (firstPos, lastPos)
    fasta = [">%s%s" % (currSeq, rng)]
    for i in xrange(0, len(ret), fastaLineLength):
        fasta.append(ret[i:i+fastaLineLength])
    return fasta

def printUsage(args, usage, minArgs=2):
  if len(args) < minArgs or args[1].startswith("-h") or args[1].startswith("--h"):
    print usage
    sys.exit(1)

usageString = """
Usage: python %s FASTA_FILE MIN_POS MAX_POS [-F]

Prints out the bases at positions [MIN_POS, MAX_POS] from
the specified FASTA file.  If -F is specified, return result
in FASTA format.
"""

if __name__=="__main__":
    printUsage(sys.argv, usageString % sys.argv[0], 4)
    if len(sys.argv)>4 and sys.argv[4]=="-F":
        fasta = readSequence(sys.argv[1], long(sys.argv[2]), long(sys.argv[3]), None, True, pad=PadSequence.NONE)
        for line in fasta: print line
    else:
        print readSequence(sys.argv[1], long(sys.argv[2]), long(sys.argv[3]))

